import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import mplfinance as mpf
from datetime import datetime
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

# 创建输出图片目录
output_dir = 'output_images'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 读取Excel文件
try:
    file_path = '../../Desktop/新建文件夹/比亚迪后复权历史行情数据.xlsx'
    df = pd.read_excel(file_path)
    
    # 任务1: 观察数据形状，输出前5行
    print("任务1: 数据基本信息")
    print(f"数据形状: {df.shape}")
    print("\n数据前5行:")
    print(df.head())
    print("\n数据列名:")
    print(df.columns.tolist())
    
    # 尝试识别日期列和价格相关列
    date_columns = []
    price_columns = []
    
    for col in df.columns:
        # 尝试识别日期列
        if '日期' in col or 'Date' in col or 'date' in col:
            date_columns.append(col)
        # 尝试识别价格相关列
        if any(keyword in col for keyword in ['开盘', '收盘', '最高', '最低', '成交量', '成交额', '涨跌幅', 'Open', 'Close', 'High', 'Low', 'Volume']):
            price_columns.append(col)
    
    print(f"\n识别的日期列: {date_columns}")
    print(f"识别的价格相关列: {price_columns}")
    
    # 假设第一个日期列是我们需要的日期
    if date_columns:
        date_col = date_columns[0]
        df[date_col] = pd.to_datetime(df[date_col])
        df.set_index(date_col, inplace=True)
    
    # 任务2: 对定量变量进行描述性分析
    print("\n任务2: 定量变量描述性分析")
    numeric_cols = df.select_dtypes(include=['float64', 'int64']).columns
    desc_stats = df[numeric_cols].describe()
    print(desc_stats)
    
    # 任务3: 计算变量之间的相关系数，绘制热力图
    print("\n任务3: 变量相关性分析")
    corr_matrix = df[numeric_cols].corr()
    print("相关系数矩阵:")
    print(corr_matrix)
    
    # 绘制相关系数热力图
    plt.figure(figsize=(12, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f', linewidths=0.5)
    plt.title('比亚迪股票数据相关系数热力图')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'correlation_heatmap.png'), dpi=300)
    plt.close()
    print("相关系数热力图已保存到output_images目录")
    
    # 任务4: 绘制2023年全年的收盘价时序图
    print("\n任务4: 绘制2023年全年收盘价时序图")
    
    # 过滤2023年的数据
    if hasattr(df.index, 'year'):
        df_2023 = df[df.index.year == 2023]
        
        # 尝试找到收盘价列
        close_col = None
        for col in df_2023.columns:
            if any(keyword in col for keyword in ['收盘', 'Close', 'close']):
                close_col = col
                break
        
        if close_col:
            plt.figure(figsize=(14, 7))
            plt.plot(df_2023.index, df_2023[close_col])
            plt.title('2023年比亚迪股票收盘价时序图')
            plt.xlabel('日期')
            plt.ylabel('收盘价')
            plt.grid(True)
            plt.tight_layout()
            plt.savefig(os.path.join(output_dir, '2023_close_price.png'), dpi=300)
            plt.close()
            print("2023年收盘价时序图已保存到output_images目录")
        else:
            print("未找到收盘价列，无法绘制时序图")
    else:
        print("数据索引不是日期类型，无法按年份过滤")
    
    # 任务5: 绘制2023年4月到6月的K线图
    print("\n任务5: 绘制2023年4月至6月K线图")
    
    # 过滤2023年4月-6月的数据
    if hasattr(df.index, 'year'):
        df_apr_jun_2023 = df[(df.index.year == 2023) & 
                            (df.index.month >= 4) & 
                            (df.index.month <= 6)]
        
        # 尝试找到OHLC列
        ohlc_cols = {'Open': None, 'High': None, 'Low': None, 'Close': None, 'Volume': None}
        
        for key in ohlc_cols.keys():
            for col in df_apr_jun_2023.columns:
                if key.lower() in col.lower() or \
                   (key == 'Open' and '开盘' in col) or \
                   (key == 'High' and '最高' in col) or \
                   (key == 'Low' and '最低' in col) or \
                   (key == 'Close' and '收盘' in col) or \
                   (key == 'Volume' and '成交量' in col):
                    ohlc_cols[key] = col
                    break
        
        # 检查是否找到了必要的列
        has_ohlc = all(val is not None for val in [ohlc_cols['Open'], ohlc_cols['High'], ohlc_cols['Low'], ohlc_cols['Close']])
        
        if has_ohlc:
            # 准备K线图数据
            kline_data = df_apr_jun_2023.rename(columns={
                ohlc_cols['Open']: 'Open',
                ohlc_cols['High']: 'High',
                ohlc_cols['Low']: 'Low',
                ohlc_cols['Close']: 'Close'
            })
            
            if ohlc_cols['Volume'] is not None:
                kline_data = kline_data.rename(columns={ohlc_cols['Volume']: 'Volume'})
                volume = True
            else:
                volume = False
            
            # 绘制K线图
            plt.figure(figsize=(14, 8))
            # 使用自定义样式来确保中文显示正常
            custom_style = mpf.make_mpf_style(base_mpf_style='charles', 
                                              rc={'font.family': ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']})
            mpf.plot(kline_data, type='candle', style=custom_style, 
                     title='2023年4月至6月比亚迪股票K线图', 
                     ylabel='价格', volume=volume, 
                     savefig=os.path.join(output_dir, '2023_apr_jun_kline.png'))
            plt.close()
            print("2023年4月至6月K线图已保存到output_images目录")
        else:
            print(f"未找到足够的OHLC列来绘制K线图。找到的列: {ohlc_cols}")
    else:
        print("数据索引不是日期类型，无法按月份过滤")
    
    print("\n所有分析任务已完成！图片已保存到output_images目录。")
    

except Exception as e:
    print(f"处理文件时出错: {e}")